package com.ybwh.springboot2.common.mybatis.plugin.result.sensitive2;

import cn.hutool.core.util.ReflectUtil;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.ibatis.binding.MapperMethod;
import org.apache.ibatis.executor.parameter.ParameterHandler;
import org.apache.ibatis.executor.resultset.ResultSetHandler;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Signature;
import org.springframework.core.annotation.AnnotationUtils;
import org.springframework.stereotype.Component;

import java.lang.reflect.Field;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.Properties;
import java.util.Set;
import java.util.stream.Collectors;

@Component
@Intercepts({@Signature(type = StatementHandler.class, method = "prepare", args = {Connection.class, Integer.class}),
        @Signature(type = ParameterHandler.class, method = "setParameters", args = PreparedStatement.class),
        @Signature(type = ResultSetHandler.class, method = "handleResultSets", args = {Statement.class})})
@Slf4j
@RequiredArgsConstructor
public class AesInterceptor implements Interceptor {

    private final AesServiceImpl aesService;

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        Object target = invocation.getTarget();
        // 拦截sql结果处理器（数据库查询结果映射到实体类）
        if (target instanceof ResultSetHandler) {
            return resultDecrypt(invocation);
        }
        // 拦截sql参数处理器（插入、更新或删除操作）
        if (target instanceof ParameterHandler) {
            return parameterEncrypt(invocation);
        }
        // 拦截sql语句处理器
        if (target instanceof StatementHandler) {
            return replaceSql(invocation);
        }
        return invocation.proceed();
    }

    @Override
    public Object plugin(Object target) {
        return Interceptor.super.plugin(target);
    }

    @Override
    public void setProperties(Properties properties) {
        Interceptor.super.setProperties(properties);
    }

    /**
     * 对mybatis映射结果进行字段解密
     *
     * @param invocation 参数
     * @return 结果
     * @throws Throwable 异常
     */
    private Object resultDecrypt(Invocation invocation) throws Throwable {
        // 取出查询的结果
        Object resultObject = invocation.proceed();
        if (Objects.isNull(resultObject)) {
            return null;
        }
        // 基于selectList
        if (resultObject instanceof ArrayList) {
            List resultList = (ArrayList) resultObject;
            if (CollectionUtils.isEmpty(resultList)) {
                return resultObject;
            }
            for (Object result : resultList) {
                if (aesService.needToDecrypt(result)) {
                    // 逐一解密
                    aesService.decrypt(result);
                }
            }
            // 基于selectOne
        } else {
            if (aesService.needToDecrypt(resultObject)) {
                aesService.decrypt(resultObject);
            }
        }
        return resultObject;
    }

    /**
     * mybatis映射参数进行加密
     *
     * @param invocation 参数
     * @return 结果
     * @throws Throwable 异常
     */
    private Object parameterEncrypt(Invocation invocation) throws Throwable {
        // @Signature 指定了 type= parameterHandler 后，这里的 invocation.getTarget() 便是parameterHandler
        ParameterHandler parameterHandler = (ParameterHandler) invocation.getTarget();
        // 获取参数对像，即 mapper 中 paramsType 的实例
        Field parameterField = parameterHandler.getClass().getDeclaredField("parameterObject");
        parameterField.setAccessible(true);
        // 取出实例
        Object parameterObject = parameterField.get(parameterHandler);

        if (parameterHandler.getParameterObject() instanceof MapperMethod.ParamMap) {
            MapperMethod.ParamMap paramMap = (MapperMethod.ParamMap) parameterHandler.getParameterObject();
            parameterObject = paramMap.get("param1");
        }
        if (null == parameterObject) {
            return invocation.proceed();
        }
        Class<?> parameterObjectClass = parameterObject.getClass();
        // 校验该实例的类是否被@SensitiveEntity所注解
        SensitiveEntity sensitiveEntity = AnnotationUtils.findAnnotation(parameterObjectClass, SensitiveEntity.class);
        // 未被@SensitiveEntity所注解 则为null，正常运行
        if (Objects.isNull(sensitiveEntity)) {
            return invocation.proceed();
        }
        // 取出当前当前类所有字段，传入加密方法
        Field[] allFields = ReflectUtil.getFields(parameterObjectClass);
        aesService.encrypt(allFields, parameterObject);
        return invocation.proceed();
    }

    /**
     * 替换mybatis Sql中的加密Key
     *
     * @param invocation 参数
     * @return 结果
     * @throws Throwable 异常
     */
    private Object replaceSql(Invocation invocation) throws Throwable {
        StatementHandler statementHandler = (StatementHandler) invocation.getTarget();
        // 获取参数对象
        Object parameterObject = statementHandler.getParameterHandler().getParameterObject();
        if (parameterObject instanceof MapperMethod.ParamMap) {
            MapperMethod.ParamMap paramMap = (MapperMethod.ParamMap) parameterObject;
            Set set = (Set) paramMap.values().stream().collect(Collectors.toSet());
            for (Object param : set) {
                if (param instanceof String) {
                    continue;
                }
                if (aesService.annotatedWithSensitive(param.getClass())) {
                    aesService.encryptParameter(param);
                }
            }
        }

        // 执行原始的 SQL 语句
        return invocation.proceed();
    }
}